import os

from langchain_core.messages import SystemMessage, HumanMessage
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate
from pydantic import BaseModel, Field

from src.ai.langchain.init_llm import get_llm

os.environ["http_proxy"] = "127.0.0.1:7890"
os.environ["https_proxy"] = "127.0.0.1:7890"
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_API_KEY"] = "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"

llm = get_llm()


class Question(BaseModel):
    content: str = Field(description="问题内容")
    answer: str = Field(description="解决方式")


json_parser = JsonOutputParser(pydantic_object=Question)

messages = """
{user_input}，\n用以下指定的JSON格式回答，格式为\n{json_structure}

"""

prompt = ChatPromptTemplate.from_template(
    messages,
    partial_variables={"json_structure": json_parser.get_format_instructions()}
)

chain = prompt | llm | json_parser

resp = chain.invoke({"user_input": "请出一道小学二年级数学题，并自己解答出来"})
print(resp)
# for s in chain.stream({"user_input": "请出一道小学二年级数学题，并自己解答出来"}):
#     print(s)
